import json
from jinja2 import Template
from scripts.server.chat_api import chat
from tqdm import tqdm

with open('./total_community_detail.json','r',encoding='utf-8') as f:
    total_community_detail = json.load(f)

prompt = '''
你是一个专业的文本分析工作者，你将基于提供的"问题"和"支撑文本"，判断该"支撑文本"能否对"问题"有较好的支撑作用。
你的回复是 0-10 的浮点数表示"支撑文本"对"问题"的支撑程度。0 表示没有支撑作用，10 表示有很好的支撑作用。请不要在回答内容中增加解释性内容。
问题：
{{question}}
支撑文本：
{{support_text}}
请根据你的判断，给出一个合理的分数。你的打分是：
'''
prompt = Template(prompt)


rag_prompt = '''
你是一个专业的文本分析工作者，具有严谨的逻辑思维和良好的语言表达能力，你将基于提供的"支撑片段"。简短且严谨的回答"问题"。
每个支撑片段都是一段文本内容的摘要，格式如下：
{{
"title": <报告标题>,
"summary": <执行摘要>,
"rating": <影响严重性评级>,
"rating_explanation": <评级解释>,
"findings": [
{{
"summary":<见解1摘要>,
"explanation": <见解1解释>
}},
{{
"summary":<见解2摘要>,
"explanation": <见解2解释>
}}
]
}}
title表示报告标题，summary对报告进行摘要。你可以参考rating的评分评估这段材料的重要程度。fingings中为该报告的一些细节，这些细节以summary(摘要)和explanation(解释)呈现。每个摘要片段以分隔符"\n-------------------\n"进行分割。


问题：
{{question}}

支撑片段：
{{support_texts}}

基于问题和支撑片段，你的回答是：
'''

if __name__== '__main__':
    query = '卫星锅的作用'
    c_level = 1  # 1-3

    cur_level = 0

    record = []
    # while cur_level > c_level:
    for i in tqdm(range(len(total_community_detail))):
        support_text = total_community_detail[i]['report']
        model_input = prompt.render(question=query,support_text=support_text)
        score = chat(model_input,[],0)
        record.append([support_text,score])

    record = sorted(record, key=lambda x: x[1], reverse=True)

    # 筛掉得分为 0 的
    input_list = [r[0] for r in record if float(r[1]) > 0]
    while True:
        try:
            support_texts = '\n-------------------\n'.join(input_list)
            model_input = rag_prompt.replace("{{question}}",query).replace("{{support_texts}}", support_texts)
            answer = chat(model_input,[],0)
            break
        except:
            input_list = input_list[:-1]
            print('长度超过，已截断')
            
    print(answer)